import re
import time

import pymysql
from tqdm import tqdm
from DataMange.t_z_code import get_data_by_date
#from LLMManage.Qwen3_8B import qwen3_8B
from LLMManage.Qwen3_8B_unsloth_bnb_4bit import qwen3_8B_4bit
#from LLMManage.zhipu import getAI
# 判断映射
JUDGMENT_MAPPING = {
    0: "【分辨不出来】",
    1: "【看涨】",
    2: "【看平】",
    3: "【看跌】",
    4: "无法匹配"
}
db_conf={"host":'localhost',
        "user":'root',
        "password":'123456',
        "database":'test03',
        "charset":'utf8mb4',}

def create_table_if_not_exists(connection, table_name):
    """创建目标表（如果不存在）"""
    sql = f"""
    CREATE TABLE IF NOT EXISTS `{table_name}` (
      `id` int NOT NULL,
      `代码` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `名称` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `阅读数` int NULL DEFAULT NULL,
      `评论数` int NULL DEFAULT NULL,
      `标题` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `内容` varchar(10000) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `作者` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `发布时间` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `回复号` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL,
      `判断` longtext CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL,
      `判断_` tinyint NULL DEFAULT NULL COMMENT '量化判断字段为数字0,1,2,3,4',
      PRIMARY KEY (`id`) USING BTREE
    ) ENGINE = InnoDB CHARACTER SET = utf8mb4 COLLATE = utf8mb4_0900_ai_ci ROW_FORMAT = DYNAMIC;
    """
    with connection.cursor() as cursor:
        cursor.execute(sql)
    connection.commit()


def get_last_processed_id(connection, table_name):
    """获取最后处理的ID"""
    with connection.cursor() as cursor:
        cursor.execute(f"SELECT MAX(id) FROM {table_name}")
        return cursor.fetchone()[0] or 0


def get_total_count(review_table,connection, last_id):
    """获取待处理总数"""
    with connection.cursor() as cursor:
        cursor.execute(f"SELECT COUNT(*) FROM {review_table} WHERE id > %s", (last_id,))
        return cursor.fetchone()[0]


def process_data(review_table,table_name, page_size=5):
    """高效安全的数据处理主函数"""
    # 建立主连接
    db_conn = pymysql.connect(**db_conf)
#
    try:
        # 创建目标表（如果不存在）
        create_table_if_not_exists(db_conn, table_name)

        # 获取最后处理的ID
        last_id = get_last_processed_id(db_conn, table_name)

        # 获取待处理总数
        total_count = get_total_count(review_table,db_conn, last_id)
        if total_count == 0:
            print("没有需要处理的新数据")
            return

        # 初始化进度条
        pbar = tqdm(total=total_count, desc="处理进度")

        # 分页处理数据
        processed = 0
        while True:
            # 每页使用独立连接和事务
            page_conn = pymysql.connect(**db_conf)

            try:
                with page_conn.cursor(pymysql.cursors.DictCursor) as cursor:
                    # 查询当前页数据
                    sql = f"""
                    SELECT * FROM {review_table} 
                    WHERE id > %s 
                    ORDER BY id 
                    LIMIT %s
                    """
                    cursor.execute(sql, (last_id, page_size))
                    page_data = cursor.fetchall()

                    if not page_data:
                        break  # 没有更多数据

                    # 处理当前页每条记录
                    for row in page_data:
                        # 生成判断
                        question = f"有网友发表帖子【{row['标题']}】他对股市的态度是看涨，看平，看跌，还是分辨不出来态度？，最终必须使用【】将其态度括起来。直接给出结果不要做任何解释！"#直接给出结果不要做任何解释！"
                        #title=row['标题']
                        #haveFourNum=re.search(r'\d{4}', title)
                        kpj,spj=get_data_by_date(row["发布时间"])
                        if kpj is not None and spj is not None:
                            question=question+f"该股票的开盘价：{kpj}，收盘价：{spj}。"
                            #print(question)
                        #startTime = time.time()
                        judgment=qwen3_8B_4bit(question)
                        #judgment = qwen3_8B(question)
                        # judgment=getAI(question)

                        #endTime = time.time()
                        #print("\n模型运行时间：",endTime-startTime)

                        # 量化判断
                        quantified = 4
                        if judgment:
                            if any(s in judgment for s in ["【分辨不出来】", "【分辨不出来态度】", "【无法分辨态度】", "【无法明确分辨其态度】"]):
                                quantified = 0
                            elif "【看涨】" in judgment:
                                quantified = 1
                            elif "【看平】" in judgment:
                                quantified = 2
                            elif "【看跌】" in judgment:
                                quantified = 3

                        # 立即插入（使用同一连接保证事务）
                        with page_conn.cursor() as insert_cursor:
                            insert_sql = f"""
                            INSERT INTO {table_name} 
                            (id, 代码, 名称, 阅读数, 评论数, 标题, 内容, 作者, 发布时间, 回复号, 判断, 判断_)
                            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                            """
                            insert_cursor.execute(insert_sql, (
                                row['id'], row['代码'], row['名称'], row['阅读数'],
                                row['评论数'], row['标题'], row['内容'], row['作者'],
                                row['发布时间'], row['回复号'], judgment, quantified
                            ))

                        # 更新最后处理的ID
                        last_id = row['id']#!
                        processed += 1
                        pbar.update(1)

                    # 提交当前页的所有插入,注意当前页数设置的合理一些，不然呈现突然停止，这一页数据就丢了
                    page_conn.commit()

            except Exception as e:
                print(f"处理过程中发生错误: {e}")
                page_conn.rollback()
                raise  # 重新抛出异常，终止处理
            finally:
                page_conn.close()

        pbar.close()
        print(f"成功处理了 {processed} 条记录")

    finally:
        db_conn.close()


if __name__ == '__main__':
    process_data(review_table="zssh_part2_part1",table_name="zssh_with_judge_qwen3_8B_zssh_part2_part1")
    #process_data(review_table="zssh", table_name="zssh_with_judge_qwen3_8B_unsloth_bnb_4bit_have_kpj_bpj")